"""
Author: yida
Time is: 2022/3/11 16:29 
this Code: 
"""
"""
Author: yida
Time is: 2022/1/24 15:13 
this Code: 转换到onnx格式
"""
import torch
import torchvision.models.detection
import torchvision.models.detection.mask_rcnn
from torchvision.models.detection.faster_rcnn import AnchorGenerator


# 加载预训练模型
def get_squeezenet1_0_model_FRCNN(num_classes):
    '''
	backbone is efficientnet
	'''
    backbone = torchvision.models.squeezenet1_0(pretrained=True).features
    # backbone = EfficientNet.from_pretrained('efficientnet-b7')
    backbone.out_channels = 512
    anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
                                       aspect_ratios=((0.5, 1.0, 2.0),))
    # anchor_generator = AnchorGenerator(sizes=((8, 16, 32, 64, 128, 256, 512),),
    # 							aspect_ratios=((0.5, 1.0, 1,5, 2.0),))
    roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
                                                    output_size=7,
                                                    sampling_ratio=2)
    model = torchvision.models.detection.faster_rcnn.FasterRCNN(backbone,
                                                                num_classes,
                                                                rpn_anchor_generator=anchor_generator,
                                                                box_roi_pool=roi_pooler)

    return model


if __name__ == '__main__':
    # with torch.no_grad():  # 在测试的时候必须加上这行和下一行代码，否则预测会出问题，这里是防止还有梯度更新这些，而且如果不加这个，后面又没有进行梯度更新的话，可能会报显存不够用的错误，我怀疑是数据没有被清理
    model = get_squeezenet1_0_model_FRCNN(12)
    model.eval()
    model = torch.load(r'/Users/yida/Desktop/31epoch/model_31.pth', map_location=torch.device('cpu'))
    start_epoch = checkpoint['epoch']
    model.load_state_dict(checkpoint['model'])

    print(model)
    model.load_state_dict(torch.load(r'I:\test\coco\coco\checkpoints/model_31.pth'),False)#再加载网络的参数
    model.load_state_dict(torch.load("/Users/yida/Desktop/31epoch/model_31.pth", map_location=torch.device('cpu')))

    # model = model.to('cpu')
    # model.eval()
    print(model)
    inputs = torch.randn(1, 1, 800, 800)
    torch.onnx.export(model, (inputs), "./faster_rcnn.onnx", verbose=True)
    print("模型转换成功!")
